
#include "glodef.h"
#include "matrices.h"
#include "errlog.h"


/****************************************************************************/


typedef FLOAT(*minfunc_t)(const FLOAT *x, unsigned n, FLOAT *grad);
static void bfgsw(MATRIX sstore, FLOAT *alpha, FLOAT *beta, unsigned ns,
        FLOAT *dsd, unsigned n, FLOAT *dnewt);
static int polyline(FLOAT *xc, FLOAT fc, FLOAT *gc, FLOAT *d, FLOAT ft,
        minfunc_t func, unsigned maxarm, unsigned n, FLOAT *xp, FLOAT *lambda);
static FLOAT polymod(FLOAT q0, FLOAT qp0, FLOAT lamc, FLOAT qc,
        FLOAT blow, FLOAT bhigh, FLOAT lamm, FLOAT qm, unsigned model);

void bfgswopt(FLOAT *x0, unsigned n, minfunc_t func)
{
    FLOAT blow, bhigh;
    unsigned numf, numg, numh;
    const unsigned maxit=20, maxarm=10, nsmax=50;
    const FLOAT tol = PERIODIC_TOL;
    unsigned itc, i, ns;
    FLOAT fc, *xi, *xt, *xc, *gc, go, *alpha, *beta, *dsd, *dsdp;
    FLOAT norm_gc, norm_dsd, dsdgc, *s, *y, old, lambda, ft;
    MATRIX sstore;
    int iarm;
    
    xc = allocV(n);
    xi = allocV(n);
    xt = allocV(n);
    y = allocV(n);
    s = allocV(n);
    gc = allocV(n);
    alpha = allocV(nsmax);
    beta = allocV(nsmax);
    dsd = allocV(n);
    dsdp = allocV(n);
    
    sstore = allocM(nsmax, n);
    
    blow=.1; bhigh=.5;
    numf=0; numg=0; numh=0;
    itc=1;
    CopyVector(x0, xc, n);
    
    fc = func(xc, n, gc);
    numf=numf+1; numg=numg+1;
    
    ns = 0;
    
    norm_gc = 0.0;
    for(i=0;i<n;i++) norm_gc += sqr(gc[i]); norm_gc = sqrt(norm_gc);
    
    /*
     * dsdp = - H_c^{-1} grad_+ if ns > 0
     */
    while(norm_gc > tol && itc <= maxit)
    {
        FLOAT alphatmp, b, b0, ys;
        
        for(i=0;i<n;i++) dsd[i] = -gc[i];
        if(ns>1)
            bfgsw(sstore, alpha, beta, ns, dsd, n, dsdp);
        else
            CopyVector(dsd, dsdp, n);
        
        /*
         * compute the direction
         */
        
        if (ns==0)
            for(i=0;i<n;i++) dsd[i] = -gc[i];
        else
        {
            FLOAT a1, a;
            for(i=0;i<n;i++) xi[i] = -dsdp[i];
            b0 = 0.0; for(i=0;i<n;i++) b0 += y[i] * s[i]; b0 = -1.0 / b0;
            
            for(a1=i=0;i<n;i++) a1 += y[i]*xi[i]; a1=-b0*(1.0-1.0/lambda)+b0*b0*a1;
            for(a=i=0;i<n;i++) a-=(a1*s[i]+b0*xi[i])*gc[i];
            /*
             * We save go=s'*g_old just so we can use it here
             * and avoid some cancellation error
             */
            alphatmp=a1+2*a/go;
            b=-b0*go;
            
            for(i=0;i<n;i++) dsd[i]=a*s[i]+b*xi[i];
        }
        
        for(norm_dsd=i=0;i<n;i++) norm_dsd += sqr(dsd[i]); norm_dsd = sqrt(norm_dsd);
        for(dsdgc=i=0;i<n;i++) dsdgc += dsd[i]*gc[i];
        if (dsdgc > -1.0e-6*norm_dsd*norm_gc)
        {
            warning_msg("bfgswopt: loss of descent");
            for(i=0;i<n;i++) dsd[i]=-gc[i];
            dsdgc = sqr(norm_gc);
            ns=0;
        }
        lambda=1.0;
        /*
         * fixup against insanely long steps see (3.50) in the book
         */
        
        lambda=fmin(1.0, 100.0/(1.0 + norm_gc));
        for(i=0;i<n;i++) xt[i]=xc[i]+lambda*dsd[i];
        ft=func(xt, n, NULL);
        numf=numf+1;
        itc=itc+1;
        
        old=1;
        for(dsdgc=i=0;i<n;i++) dsdgc += dsd[i]*gc[i];
        if (old==0)
        {
            FLOAT goal=fc+1.e-4*dsdgc;
            iarm=0;
            if (ft > goal)
            {
                iarm = polyline(xc, fc, gc, dsd, ft, func, maxarm, n, xt, &lambda);
                if(iarm == -1)
                {
                    warning_msg("%s: line search failure.\n", __func__);
                    goto clean_and_go;
                }
            }
        }
        if(old==1)
        {
            FLOAT goalval, q0, qp0, qm, lamc, lamm, qc;
            iarm=0;
            goalval=.0001*dsdgc;
            q0=fc;
            qp0=dsdgc;
            lamc=lambda;
            qc=ft;
            while(ft > fc + lambda*goalval )
            {
                iarm=iarm+1;
                if (iarm==1)
                    lambda=polymod(q0, qp0, lamc, qc, blow, bhigh, 0, 0, 2);
                else
                    lambda=polymod(q0, qp0, lamc, qc, blow, bhigh, lamm, qm, 3);
                qm=qc; lamm=lamc; lamc=lambda;
                for(i=0;i<n;i++) xt[i]=xc[i]+lambda*dsd[i];
                ft=func(xt, n, NULL);
                qc=ft;
                numf=numf+1;
                if(iarm > (int) maxarm)
                {
                    warning_msg("%s: too many backtracks in BFGS line search.\n", __func__);
                    goto clean_and_go;
                }
            }
        }
        for(i=0;i<n;i++) s[i]=xt[i]-xc[i];
        CopyVector(gc, y, n);
        for(go=i=0;i<n;i++) go+=s[i]*gc[i];
        /* lambda=norm(s)/norm(dsd); */
        CopyVector(xt, xc, n);
        fc=func(xc, n, gc);
        for(norm_gc=i=0;i<n;i++) norm_gc += gc[i]*gc[i]; norm_gc = sqrt(norm_gc);
        for(i=0;i<n;i++) y[i] = gc[i]-y[i];
        numf=numf+1; numg=numg+1;
        /*
         * restart if y'*s is not positive or we're out of room
         */
        for(ys=i=0;i<n;i++) ys += y[i]*s[i];
        if (ys <= 0 || ns==nsmax)
        {
            warning_msg("%s: loss of positivity or storage", __func__);
            ns=0;
        }
        else
        {
            ns=ns+1;
            CopyVector(s, sstore[ns-1], n);
            if(ns>1)
            {
                alpha[ns-2]=alphatmp;
                beta[ns-2]=b0/(b*lambda);
            }
        }
    }
    clean_and_go:
    CopyVector(xc, x0, n);
    
    
}


static void bfgsw(MATRIX sstore, FLOAT *alpha, FLOAT *beta, unsigned ns, FLOAT *dsd,
        unsigned n, FLOAT *dnewt)
{
    FLOAT *sigma, *gamma1, *gamma2, *gamma3, *delta;
    unsigned i, j;
    
    CopyVector(dsd, dnewt, n);
    if (ns<=1) return;
    
    sigma = allocV(ns);
    gamma1 = allocV(ns-1);
    gamma2 = allocV(ns-1);
    gamma3 = allocV(ns-1);
    delta = allocV(ns-2);
    
    
    for(j=0;j<ns;j++)
        for(i=0;i<n;i++)
            sigma[j] += sstore[j][i]*dsd[i];
    
    for(j=0;j<ns-1;j++)
    {
        gamma1[j] = sigma[j] * alpha[j];
        gamma2[j] = sigma[j] * beta[j];
        gamma3[j] = gamma1[j]+beta[j]*sigma[j+1];
    }
    
    for(i=0;i<n;i++)
        dnewt[i] += gamma3[0]*sstore[0][i] + gamma2[ns-2]*sstore[ns-1][i];
    
    if(ns > 2)
    {
        for(j=0;j<ns-2;j++)
            delta[j] = gamma2[j]+gamma3[j+1];
        
        for(i=0;i<n;i++)
            for(j=0;j<ns-2;j++)
                dnewt[i] += sstore[j+1][i]*delta[j];
    }
    
    freeV(&sigma);
    freeV(&gamma1);
    freeV(&gamma2);
    freeV(&gamma3);
    freeV(&delta);
    
    return;
}


static int polyline(FLOAT *xc, FLOAT fc, FLOAT *gc, FLOAT *d, FLOAT ft,
        minfunc_t func, unsigned maxarm, unsigned n, FLOAT *xp, FLOAT *lambda)
{
    FLOAT alp=1.e-4, blow=.1, bhigh=.5;
    FLOAT q0, qc, qp0, lamc, lamm, qm, fgoal;
    unsigned i, iarm;
    /*
     * Set up the search
     */
    q0=fc;
    for(qp0=i=0;i<n;i++) qp0 += gc[i]*d[i];
    qc=ft;
    lamc=1;
    iarm=0;
    fgoal=fc+alp*lamc*qp0;
    while (ft > fgoal)
    {
        iarm=iarm+1;
        if (iarm==1)
            (*lambda)=polymod(q0, qp0, lamc, qc, blow, bhigh, lamm, qm, 2);
        else
            (*lambda)=polymod(q0, qp0, lamc, qc, blow, bhigh, lamm, qm, 3);
        qm=qc;
        lamm=lamc;
        lamc=(*lambda);
        for(i=0;i<n;i++) xp[i]=xc[i]+(*lambda)*d[i];
        ft=func(xp, n, NULL);
        qc=ft;
        if(iarm > maxarm)
        {
            warning_msg("%s: line search failure.\n", __func__);
            CopyVector(xc, xp, n);
            return -1;
        }
        fgoal=fc+alp*lamc*qp0;
    }
    return 1;
}

static FLOAT polymod(FLOAT q0, FLOAT qp0, FLOAT lamc, FLOAT qc,
        FLOAT blow, FLOAT bhigh, FLOAT lamm, FLOAT qm, unsigned model)
/*
 * Finds minimizer lambda of the cubic polynomial q on the interval
 * [blow * lamc, bhigh * lamc] such that
 *
 * q(0) = q0, q'(0) = qp0, q(lamc) = qc, q(lamm) = qm
 *
 * if data for a cubic is not available (first stepsize reduction) then
 * q is the quadratic such that
 *
 * q(0) = q0, q'(0) = qp0, q(lamc) = qc
 */
{
    FLOAT lleft, lright, lplus;
    lleft=lamc*blow; lright=lamc*bhigh;
    if (model == 2)
        /*
         * quadratic model (temp hedge in case lamc is not 1)
         */
    {
        lplus = - qp0/(2 * lamc*(qc - q0 - qp0) );
        if (lplus < lleft) lplus = lleft;
        if (lplus > lright) lplus = lright;
    }
    else if (model == 3)
        /*
         * cubic model
         */
    {
        FLOAT _a[4], *a[2] = {_a, _a+2}, b[2], c[2];
        a[0][0]=sqr(lamc);
        a[0][1]=a[0][0]*lamc;
        a[1][0]=sqr(lamm);
        a[1][1]=a[1][0]*lamm;
        b[0] = qc - (q0+qp0*lamc);
        b[1] = qm - (q0+qp0*lamm);
        SolveLinearSystem(a, 2, b, c);
        lplus=(-c[0]+sqrt(c[0]*c[0] - 3 *c[1] *qp0))/(3*c[1]);
        if (lplus < lleft) lplus = lleft;
        if (lplus > lright) lplus = lright;
    }
    else
        error_msg("%s: model must be 2 or 3", __func__);
    return lplus;
}



